package graph

import "gitee.com/yrwy/msgo/pkg/mem/stl/list"

type Searcher[T comparable] interface {
	Adjacency(prev T) []T
	NewTree() Tree[T]
}

func BFS[T comparable](x Searcher[T], s T) Tree[T] {
	return BreadthFirstSearch[T](x, s, -1, defaultDegree[T])
}

func DFS[T comparable](x Searcher[T], s T) Tree[T] {
	return DepthFirstSearch[T](x, s, -1, defaultDegree[T])
}

// u: 前节点
// v: 后节点
type DegreeFunc[T any] func(u, v T) int

func defaultDegree[T any](T, T) int { return 1 }

// s : 起始位置
// max_degree : 限制深度
func BreadthFirstSearch[T comparable](x Searcher[T], s T, max_degree int, degree DegreeFunc[T]) Tree[T] {
	df := degree
	if df == nil {
		df = defaultDegree[T]
	}
	d := make(map[T]int) //度
	p := x.NewTree()     //前节点
	l := list.New[T](nil)
	d[s] = 0
	l.PushBack(s)
	for l.Len() > 0 {
		u := l.PopFront()
		u_degree := d[u]
		adj := x.Adjacency(u)
		for _, v := range adj {
			v_d := u_degree + df(u, v)
			if nd, ok := d[v]; !ok {
				if max_degree < 0 || v_d <= max_degree {
					l.PushBack(v)
					d[v] = v_d
					p.Set(v, u)
				}
			} else if max_degree >= 0 && v_d < nd {
				d[v] = v_d
				p.Set(v, u)
				l.PushBack(v)
			}
		}
	}
	return p
}

func DepthFirstSearch[T comparable](x Searcher[T], s T, max_degree int, degree DegreeFunc[T]) Tree[T] {
	d := newDfs(x, s)
	df := degree
	if df == nil {
		df = defaultDegree[T]
	}
	d.search(s, max_degree, degree)
	return d.p
}

type dfs[T comparable] struct {
	r Searcher[T]
	s T
	p Tree[T]   //返回深度优先森林
	d map[T]int //度
}

func newDfs[T comparable](r Searcher[T], s T) *dfs[T] {
	d := dfs[T]{r: r, s: s}
	d.d = make(map[T]int) //度
	d.p = r.NewTree()
	d.d[s] = 0
	return &d
}

func (x *dfs[T]) search(u T, max_degree int, degree DegreeFunc[T]) {
	u_degree := x.d[u]
	adj := x.r.Adjacency(u)
	for _, v := range adj {
		d := u_degree + degree(u, v)
		if nd, ok := x.d[v]; !ok {
			if max_degree < 0 || d <= max_degree { //超过了du则走其他的路径
				x.p.Set(v, u)
				x.d[v] = d
				x.search(v, max_degree, degree)
			}
		} else if max_degree >= 0 && d < nd {
			x.p.Set(v, u)
			x.d[v] = d
			x.search(v, max_degree, degree)
		}
	}
}
